import torch
from utils import check_state_dict
from models.birefnet import BiRefNet
from torchvision.ops.deform_conv import DeformConv2d
import deform_conv2d_onnx_exporter
import argparse
from safetensors.torch import load_file

def parse_tuple(input_str):
    try:
        return tuple(map(int, input_str.split(',')))
    except:
        raise argparse.ArgumentTypeError("--input_shape元组格式错误！请使用 'height,width' 格式（例如：'1024,1024'）")

def parse_arguments():
    parser = argparse.ArgumentParser(description='export onnx weight of birefnet')
    parser.add_argument('--weight', type=str, default="BiRefNet-general-epoch_244.pth", help='pytorch pth weight')
    parser.add_argument('--onnx_name', type=str, default="BiRefNet-general-epoch_244", help='')
    parser.add_argument('--weight_type', type=str, default="pth",help='input shape')
    parser.add_argument('--is_dynamic', action="store_true",help='whether to export dynamic shape onnx')
    parser.add_argument(
        '--input_shape',
        type=parse_tuple,
        default=(1024, 1024),
        help='input resolution'
    )
    args = parser.parse_args()
    return args

def main(args):
    birefnet = BiRefNet(bb_pretrained=False)
    if args.weight_type == "pth":
        state_dict = torch.load(args.weight, map_location='cpu')
        state_dict = check_state_dict(state_dict)
        birefnet.load_state_dict(state_dict)
    elif args.weight_type == "safetensors":
        loaded_state_dict=load_file(args.weight)
        birefnet.load_state_dict(loaded_state_dict)
    else:
        ValueError("weight format only supports pth or safetensors")

    torch.set_float32_matmul_precision(['high', 'highest'][0])
    birefnet.to('cpu')
    birefnet.eval()

    # deform_conv2d op登録
    deform_conv2d_onnx_exporter.register_deform_conv2d_onnx_op()

    # def convert_to_onnx(net, file_name='output.onnx', input_shape=(1024, 1024), device='cpu'):
    input_shape=args.input_shape
    file_name=args.onnx_name + ".onnx"
    input = torch.randn(1, 3, input_shape[0], input_shape[1]).to('cpu')

    input_layer_names = ['input_image']
    output_layer_names = ['output_image']

    if args.is_dynamic:
        print("--------------start to export dynamic shape onnx ------------")
        torch.onnx.export(
        birefnet,
        input,
        file_name,
        verbose=True,
        opset_version=12,
        input_names=input_layer_names,
        output_names=output_layer_names,
        dynamic_axes={'input_image': {2: 'input_height', 3: 'input_width'}})
    else:
        torch.onnx.export(
        birefnet,
        input,
        file_name,
        verbose=True,
        opset_version=12,
        input_names=input_layer_names,
        output_names=output_layer_names,
        )
    print("--------------expor dynamic shape onnx successfully! ------------")
    
if __name__=="__main__":
    args=parse_arguments()
    main(args)